import os
import re
import shutil
import time
from dataclasses import dataclass

import cv2
import base64
import numpy as np
from typing import List, Optional

import pandas as pd
from fastapi import FastAPI, File, UploadFile, HTTPException, Form, Body
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from marshmallow.fields import String
from paddlex import create_pipeline
from pydantic import BaseModel

from src.crop_nutrition_content import crop_nutrition_content
from src.ocr_processor import OCRProcessor
from src.preprocess_nutrition_table import preprocess_grid


@dataclass
class NutritionItem:
    id: int
    name: str
    unit: str

# 数据库模拟
NUTRITION_DB = {
    "能量": {"id": 1, "unit": "kcal"},
    "蛋白质": {"id": 2, "unit": "g"},
    "脂肪": {"id": 3, "unit": "g"},
    "碳水化合物": {"id": 4, "unit": "g"},
    "纤维": {"id": 5, "unit": "g"},
    "糖": {"id": 6, "unit": "g"},
    "钠": {"id": 7, "unit": "mg"},
    "钾": {"id": 8, "unit": "mg"},
    "钙": {"id": 9, "unit": "mg"},
    "铁": {"id": 10, "unit": "mg"},
    "维生素C": {"id": 11, "unit": "mg"},
    "维生素A": {"id": 12, "unit": "IU"},
    "胆固醇": {"id": 13, "unit": "mg"},
    "饱和脂肪": {"id": 14, "unit": "g"},
    "单不饱和脂肪": {"id": 15, "unit": "g"},
    "多不饱和脂肪": {"id": 16, "unit": "g"},
    "维生素D": {"id": 17, "unit": "IU"},
}

# 转换为列表，方便比对
NUTRITION_ITEMS = [
    NutritionItem(id=value["id"], name=key, unit=value["unit"])
    for key, value in NUTRITION_DB.items()
]
# 创建应用
app = FastAPI(
    title="PaddleOCR API",
    description="OCR识别API服务，基于PaddleOCR提供文本识别功能",
    version="1.0.0"
)

# 配置CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 允许所有来源，生产环境应该设置为特定域名
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 配置静态文件服务
app.mount("/static", StaticFiles(directory="static"), name="static")

# 配置上传目录
UPLOAD_FOLDER = "static/uploads"
ALLOWED_EXTENSIONS = {"png", "jpg", "jpeg"}
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

# 初始化OCR处理器
ocr_processor = OCRProcessor()
# 初始化表格识别pipeline
table_pipeline = create_pipeline(pipeline="table_recognition_v2")

# 定义数据模型
class OCRResult(BaseModel):
    text: str
    confidence: float
    box: Optional[List[List[int]]] = None

class OCRResponse(BaseModel):
    results: List[OCRResult]
    original_image_url: Optional[str] = None
    annotated_image_url: Optional[str] = None
    annotated_image_base64: Optional[str] = None

class Base64Request(BaseModel):
    image_base64: str
    return_annotated: bool = True

class RecognizedNutritionItem(BaseModel):
    nutritionId: int
    nutritionName: str
    unit: Optional[str] = None
    amount: Optional[float] = None
    isMatched: bool
    ocrName: Optional[str] = None

class TableData(BaseModel):
    headers: List[str]
    rows: List[List[str]]

class RecognizedNutritionResponse(BaseModel):
    debugTables: List[TableData]  # 结构化表格数据
    debugRawPairs: List[dict]     # 原始键值对
    debugPreprocessedPath: str    # 预处理图片路径
    results: List[RecognizedNutritionItem]

@app.get("/")
async def root():
    """API根路径，返回简单信息"""
    return {"message": "欢迎使用PaddleOCR API服务", "status": "运行中"}

@app.post("/api/recognize", response_model=OCRResponse)
async def recognize_from_upload(
    file: UploadFile = File(...),
    return_annotated: bool = Form(True)
):
    """
    从上传文件中识别文本
    
    - **file**: 上传的图像文件
    - **return_annotated**: 是否返回标注后的图像
    """
    # 验证文件类型
    file_ext = file.filename.split(".")[-1].lower()
    if file_ext not in ALLOWED_EXTENSIONS:
        raise HTTPException(status_code=400, detail="不支持的文件类型")
    
    # 保存上传文件
    file_path = os.path.join(UPLOAD_FOLDER, file.filename)
    with open(file_path, "wb") as buffer:
        buffer.write(await file.read())
    
    try:
        # OCR识别
        ocr_results, annotated_image = ocr_processor.recognize_text(file_path)
        
        # 准备返回数据
        results = [
            OCRResult(
                text=item["text"],
                confidence=item["confidence"],
                box=item.get("box")
            ) for item in ocr_results
        ]
        
        # 构建响应
        response = {"results": results}
        
        # 添加图像URL
        response["original_image_url"] = f"/static/uploads/{file.filename}"
        
        # 如果需要返回标注图像
        if return_annotated:
            # 保存标注图像
            annotated_filename = f"result_{file.filename}"
            annotated_path = os.path.join(UPLOAD_FOLDER, annotated_filename)
            cv2.imwrite(annotated_path, annotated_image)
            
            # 添加标注图像URL
            response["annotated_image_url"] = f"/static/uploads/{annotated_filename}"
            
            # 转换为base64格式方便移动端直接使用
            _, buffer = cv2.imencode(f".{file_ext}", annotated_image)
            img_base64 = base64.b64encode(buffer).decode("utf-8")
            response["annotated_image_base64"] = img_base64
        
        return response
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"处理失败: {str(e)}")

@app.post("/api/recognize/base64", response_model=OCRResponse)
async def recognize_from_base64(request: Base64Request):
    """
    从Base64编码的图像中识别文本
    
    - **image_base64**: Base64编码的图像
    - **return_annotated**: 是否返回标注后的图像
    """
    try:
        # 解码Base64图像
        image_data = base64.b64decode(request.image_base64)
        nparr = np.frombuffer(image_data, np.uint8)
        image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        
        # 保存为临时文件
        temp_filename = f"temp_{os.urandom(8).hex()}.jpg"
        temp_path = os.path.join(UPLOAD_FOLDER, temp_filename)
        cv2.imwrite(temp_path, image)
        
        # OCR识别
        ocr_results, annotated_image = ocr_processor.recognize_text(temp_path)
        
        # 准备返回数据
        results = [
            OCRResult(
                text=item["text"],
                confidence=item["confidence"],
                box=item.get("box")
            ) for item in ocr_results
        ]
        
        # 构建响应
        response = {"results": results}
        
        # 如果需要返回标注图像
        if request.return_annotated:
            # 转换为base64格式
            _, buffer = cv2.imencode(".jpg", annotated_image)
            img_base64 = base64.b64encode(buffer).decode("utf-8")
            response["annotated_image_base64"] = img_base64
        
        # 删除临时文件
        os.remove(temp_path)
        
        return response
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"处理失败: {str(e)}")

@app.post("/api/preprocess")
async def preprocess_image(file: UploadFile = File(...)):
    """
    对上传的图像进行预处理，以提高OCR识别率
    
    - **file**: 上传的图像文件
    """
    # 验证文件类型
    file_ext = file.filename.split(".")[-1].lower()
    if file_ext not in ALLOWED_EXTENSIONS:
        raise HTTPException(status_code=400, detail="不支持的文件类型")
    
    # 保存上传文件
    file_path = os.path.join(UPLOAD_FOLDER, file.filename)
    with open(file_path, "wb") as buffer:
        buffer.write(await file.read())
    
    try:
        # 图像预处理
        processed_path = ocr_processor.preprocess_image(file_path)
        processed_filename = os.path.basename(processed_path)
        
        # 返回处理后的图像URL
        return {
            "original_image_url": f"/static/uploads/{file.filename}",
            "processed_image_url": f"/static/uploads/{processed_filename}"
        }
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"处理失败: {str(e)}")


@app.post("/api/recognize/nutrition-table", response_model=RecognizedNutritionResponse)
async def recognize_nutrition_table(file: UploadFile = File(...)):
    try:

        # 保存上传文件
        filename = f"nutrition_{int(time.time())}_{file.filename}"
        file_path = os.path.join(UPLOAD_FOLDER, filename)
        with open(file_path, "wb") as buffer:
            buffer.write(await file.read())

        # 进行图像预处理（示例函数，需要实际实现）
        cropped_path = crop_nutrition_content(file_path)
        gridded_path = preprocess_grid(cropped_path)
        processed_filename = os.path.basename(gridded_path)

        # 表格识别
        output = table_pipeline.predict(
            input=gridded_path,
            use_doc_orientation_classify=True,
            use_doc_unwarping=False,
            return_ocr_result_in_table=True,
            use_layout_detection = False,
        )

        # 解析表格数据
        tables = []
        raw_pairs = []
        results = []

        for table_result in output:
            # 获取结构化表格数据
            print(">>> Raw Table Result:", table_result)
            print(">>> Keys in __dict__:", table_result.__dict__.keys())
            html_dict = table_result._to_html()
            for table_html in html_dict.values():
                try:
                    df_list = pd.read_html(table_html)
                    if not df_list:
                        continue
                    df = df_list[0]

                    # 构建表格数据结构
                    tables.append(TableData(
                        headers=[str(col) for col in df.columns.tolist()],
                        rows=[[str(cell) if not pd.isna(cell) else "" for cell in row] for row in df.values.tolist()]
                    ))

                    # 提取原始键值对（假设第一列为名称，第二列为值）
                    for _, row in df.iterrows():
                        if len(row) >= 2:
                            raw_pairs.append({
                                "name": str(row[0]).strip(),
                                "value": str(row[1]).strip()
                            })

                except Exception as e:
                    print(f"表格解析错误: {e}")

            # 营养匹配逻辑
            for table in tables:
                for row in table.rows:
                    nutrition_name = None
                    amount = None
                    unit = None

                    # 查找营养名称
                    for cell in row:
                        cell_text = str(cell).strip()
                        for nutrition in NUTRITION_ITEMS:
                            if nutrition.name in cell_text:
                                nutrition_name = nutrition.name
                                break
                        if nutrition_name:
                            break

                    # 查找数值和单位
                    if nutrition_name:
                        for cell in row:
                            cell_text = str(cell).strip()
                            # 匹配数字和单位（增强正则表达式）
                            match = re.match(
                                r"([\d,]+\.?\d*)\s*([a-zA-Z%°]+)?",
                                cell_text
                            )
                            if match:
                                amount_str = match.group(1).replace(",", "")
                                amount = float(amount_str)
                                unit = match.group(2) if match.group(2) else ""
                                break

                        # 匹配数据库
                        db_item = NUTRITION_DB.get(nutrition_name)
                        is_matched = bool(db_item)
                        # 安全获取字段
                        nutrition_id = db_item["id"] if is_matched else -1
                        default_unit = db_item["unit"] if is_matched else ""

                        results.append(RecognizedNutritionItem(
                            nutritionId=nutrition_id,
                            nutritionName=nutrition_name if is_matched else "未知",
                            unit=unit or default_unit,
                            amount=amount,
                            isMatched=is_matched,
                            ocrName=nutrition_name
                        ))

        return {
            "debugTables": tables,  # 转换为Pydantic模型
            "debugRawPairs": raw_pairs,
            "debugPreprocessedPath": f"/static/preprocess/grid/{processed_filename}",
            "results": results,  # 确保results是合法数据
            "weight": 100
        }

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))